#ifndef __MODULE__
#define __MODULE__

#include <vector>
#include <memory>
#include <iostream>
#include "../Tensor.hpp"

template<typename _T = float>
class Module {
public:
    Module() {}

    // Virtual destructor to ensure proper cleanup of derived classes
    virtual ~Module() {}

    // Pure virtual forward function that must be overridden by all subclasses
    virtual DimX::Tensor<_T> forward(DimX::Tensor<_T> x) {
        return DimX::Tensor<_T>();
    };

    // Add a method to register parameters if needed
    void register_parameter(const std::string& name, std::shared_ptr<double> parameter) {
        parameters.push_back(parameter);
        parameter_names.push_back(name);
    }

    // Optionally, add methods for handling submodules if your design requires
    void add_module(const std::string& name, std::shared_ptr<Module> module) {
        modules.push_back(module);
        module_names.push_back(name);
    }

private:
    std::vector<std::shared_ptr<double>> parameters;
    std::vector<std::string> parameter_names;
    std::vector<std::shared_ptr<Module>> modules;
    std::vector<std::string> module_names;

    // Base class for all neural network modules.
    // Your models should also subclass this class.
    // Modules can also contain other Modules, allowing to nest them in a tree structure.
};

#endif // __MODULE__